{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5a6f7481-4064-4e8c-9b94-21905d6db6d7",
   "metadata": {},
   "source": [
    "# Lesson 4: Technical Debt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e19aa2e-8c1b-43ff-b540-cea603a79b5c",
   "metadata": {},
   "source": [
    "#### Setup\n",
    "Set the MakerSuite API key with the provided helper function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13f21190-1c61-4fb8-957d-ef5d070bf333",
   "metadata": {
    "height": 64
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from utils import get_api_key\n",
    "import google.generativeai as palm\n",
    "from google.api_core import client_options as client_options_lib\n",
    "\n",
    "palm.configure(\n",
    "    api_key=get_api_key(),\n",
    "    transport=\"rest\",\n",
    "    client_options=client_options_lib.ClientOptions(\n",
    "        api_endpoint=os.getenv(\"GOOGLE_API_BASE\"),\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b7550f9-c19b-4bdf-83d7-61ee51d94083",
   "metadata": {},
   "source": [
    "#### Pick the model that generates text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eff6eccb-49ba-439c-9c11-94d7d9dc8902",
   "metadata": {
    "height": 64
   },
   "outputs": [],
   "source": [
    "models = [m for m in palm.list_models() if 'generateText' in m.supported_generation_methods]\n",
    "model_bison = models[0]\n",
    "model_bison"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "089429a3-357d-4df2-a11a-9ba4bc9b7e85",
   "metadata": {},
   "source": [
    "#### Helper function to call the PaLM API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87831e1c-493f-473c-a13a-f7856678aef8",
   "metadata": {
    "height": 149
   },
   "outputs": [],
   "source": [
    "from google.api_core import retry\n",
    "@retry.Retry()\n",
    "def generate_text(prompt, \n",
    "                  model=model_bison, \n",
    "                  temperature=0.0):\n",
    "    return palm.generate_text(prompt=prompt,\n",
    "                              model=model,\n",
    "                              temperature=temperature)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87eef99a-7840-4296-beaf-ea0b758bb9a1",
   "metadata": {},
   "source": [
    "### Ask an LLM to explain a complex code base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71c2070a-56f6-4c88-b5eb-f6cb0ca5e6a7",
   "metadata": {
    "height": 4382
   },
   "outputs": [],
   "source": [
    "#@title Complex Code Block\n",
    "# Note: Taken from https://github.com/lmoroney/odmlbook/blob/63c0825094b2f44efc5c4d3226425a51990e73d6/BookSource/Chapter08/ios/cats_vs_dogs/CatVsDogClassifierSample/ModelDataHandler/ModelDataHandler.swift\n",
    "CODE_BLOCK = \"\"\"\n",
    "// Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n",
    "//\n",
    "// Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "// you may not use this file except in compliance with the License.\n",
    "// You may obtain a copy of the License at\n",
    "//\n",
    "//    http://www.apache.org/licenses/LICENSE-2.0\n",
    "//\n",
    "// Unless required by applicable law or agreed to in writing, software\n",
    "// distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "// See the License for the specific language governing permissions and\n",
    "// limitations under the License.\n",
    "\n",
    "import CoreImage\n",
    "import TensorFlowLite\n",
    "import UIKit\n",
    "\n",
    "\n",
    "/// An inference from invoking the `Interpreter`.\n",
    "struct Inference {\n",
    "  let confidence: Float\n",
    "  let label: String\n",
    "}\n",
    "\n",
    "/// Information about a model file or labels file.\n",
    "typealias FileInfo = (name: String, extension: String)\n",
    "\n",
    "/// Information about the MobileNet model.\n",
    "enum MobileNet {\n",
    "  static let modelInfo: FileInfo = (name: \"converted_model\", extension: \"tflite\")\n",
    "}\n",
    "\n",
    "/// This class handles all data preprocessing and makes calls to run inference on a given frame\n",
    "/// by invoking the `Interpreter`. It then formats the inferences obtained and returns the top N\n",
    "/// results for a successful inference.\n",
    "class ModelDataHandler {\n",
    "\n",
    "  // MARK: - Public Properties\n",
    "\n",
    "  /// The current thread count used by the TensorFlow Lite Interpreter.\n",
    "  let threadCount: Int\n",
    "\n",
    "  let resultCount = 1\n",
    "\n",
    "  // MARK: - Model Parameters\n",
    "\n",
    "  let batchSize = 1\n",
    "  let inputChannels = 3\n",
    "  let inputWidth = 224\n",
    "  let inputHeight = 224\n",
    "\n",
    "  // MARK: - Private Properties\n",
    "\n",
    "  /// List of labels from the given labels file.\n",
    "  private var labels: [String] = [\"Cat\", \"Dog\"]\n",
    "\n",
    "  /// TensorFlow Lite `Interpreter` object for performing inference on a given model.\n",
    "  private var interpreter: Interpreter\n",
    "\n",
    "  /// Information about the alpha component in RGBA data.\n",
    "  private let alphaComponent = (baseOffset: 4, moduloRemainder: 3)\n",
    "\n",
    "  // MARK: - Initialization\n",
    "\n",
    "  /// A failable initializer for `ModelDataHandler`. A new instance is created if the model and\n",
    "  /// labels files are successfully loaded from the app's main bundle. Default `threadCount` is 1.\n",
    "  init?(modelFileInfo: FileInfo, threadCount: Int = 1) {\n",
    "    let modelFilename = modelFileInfo.name\n",
    "\n",
    "    // Construct the path to the model file.\n",
    "    guard let modelPath = Bundle.main.path(\n",
    "      forResource: modelFilename,\n",
    "      ofType: modelFileInfo.extension\n",
    "      ) else {\n",
    "        print(\"Failed to load the model file with name: \\(modelFilename).\")\n",
    "        return nil\n",
    "    }\n",
    "\n",
    "    // Specify the options for the `Interpreter`.\n",
    "    self.threadCount = threadCount\n",
    "    var options = InterpreterOptions()\n",
    "    options.threadCount = threadCount\n",
    "    do {\n",
    "      // Create the `Interpreter`.\n",
    "      interpreter = try Interpreter(modelPath: modelPath, options: options)\n",
    "    } catch let error {\n",
    "      print(\"Failed to create the interpreter with error: \\(error.localizedDescription)\")\n",
    "      return nil\n",
    "    }\n",
    "\n",
    "  }\n",
    "\n",
    "  // MARK: - Public Methods\n",
    "\n",
    "  /// Performs image preprocessing, invokes the `Interpreter`, and process the inference results.\n",
    "  func runModel(onFrame pixelBuffer: CVPixelBuffer) -> [Inference]? {\n",
    "    let sourcePixelFormat = CVPixelBufferGetPixelFormatType(pixelBuffer)\n",
    "    assert(sourcePixelFormat == kCVPixelFormatType_32ARGB ||\n",
    "      sourcePixelFormat == kCVPixelFormatType_32BGRA ||\n",
    "      sourcePixelFormat == kCVPixelFormatType_32RGBA)\n",
    "\n",
    "\n",
    "    let imageChannels = 4\n",
    "    assert(imageChannels >= inputChannels)\n",
    "\n",
    "    // Crops the image to the biggest square in the center and scales it down to model dimensions.\n",
    "    let scaledSize = CGSize(width: inputWidth, height: inputHeight)\n",
    "    guard let thumbnailPixelBuffer = pixelBuffer.centerThumbnail(ofSize: scaledSize) else {\n",
    "      return nil\n",
    "    }\n",
    "\n",
    "    let outputTensor: Tensor\n",
    "    do {\n",
    "      // Allocate memory for the model's input `Tensor`s.\n",
    "      try interpreter.allocateTensors()\n",
    "\n",
    "      // Remove the alpha component from the image buffer to get the RGB data.\n",
    "      guard let rgbData = rgbDataFromBuffer(\n",
    "        thumbnailPixelBuffer,\n",
    "        byteCount: batchSize * inputWidth * inputHeight * inputChannels\n",
    "        ) else {\n",
    "          print(\"Failed to convert the image buffer to RGB data.\")\n",
    "          return nil\n",
    "      }\n",
    "\n",
    "      // Copy the RGB data to the input `Tensor`.\n",
    "      try interpreter.copy(rgbData, toInputAt: 0)\n",
    "\n",
    "      // Run inference by invoking the `Interpreter`.\n",
    "      try interpreter.invoke()\n",
    "\n",
    "      // Get the output `Tensor` to process the inference results.\n",
    "      outputTensor = try interpreter.output(at: 0)\n",
    "    } catch let error {\n",
    "      print(\"Failed to invoke the interpreter with error: \\(error.localizedDescription)\")\n",
    "      return nil\n",
    "    }\n",
    "\n",
    "    let results = [Float32](unsafeData: outputTensor.data) ?? []\n",
    "\n",
    "    // Process the results.\n",
    "    let topNInferences = getTopN(results: results)\n",
    "\n",
    "    // Return the inference time and inference results.\n",
    "    return topNInferences\n",
    "  }\n",
    "\n",
    "  // MARK: - Private Methods\n",
    "\n",
    "  /// Returns the top N inference results sorted in descending order.\n",
    "  private func getTopN(results: [Float]) -> [Inference] {\n",
    "    // Create a zipped array of tuples [(labelIndex: Int, confidence: Float)].\n",
    "    let zippedResults = zip(labels.indices, results)\n",
    "\n",
    "    // Sort the zipped results by confidence value in descending order.\n",
    "    let sortedResults = zippedResults.sorted { $0.1 > $1.1 }.prefix(resultCount)\n",
    "\n",
    "    // Return the `Inference` results.\n",
    "    return sortedResults.map { result in Inference(confidence: result.1, label: labels[result.0]) }\n",
    "  }\n",
    "\n",
    "  /// Loads the labels from the labels file and stores them in the `labels` property.\n",
    "  private func loadLabels(fileInfo: FileInfo) {\n",
    "    let filename = fileInfo.name\n",
    "    let fileExtension = fileInfo.extension\n",
    "    guard let fileURL = Bundle.main.url(forResource: filename, withExtension: fileExtension) else {\n",
    "      fatalError(\"Labels file not found in bundle. Please add a labels file with name \" +\n",
    "        \"\\(filename).\\(fileExtension) and try again.\")\n",
    "    }\n",
    "    do {\n",
    "      let contents = try String(contentsOf: fileURL, encoding: .utf8)\n",
    "      labels = contents.components(separatedBy: .newlines)\n",
    "    } catch {\n",
    "      fatalError(\"Labels file named \\(filename).\\(fileExtension) cannot be read. Please add a \" +\n",
    "        \"valid labels file and try again.\")\n",
    "    }\n",
    "  }\n",
    "\n",
    "  /// Returns the RGB data representation of the given image buffer with the specified `byteCount`.\n",
    "  ///\n",
    "  /// - Parameters\n",
    "  ///   - buffer: The pixel buffer to convert to RGB data.\n",
    "  ///   - byteCount: The expected byte count for the RGB data calculated using the values that the\n",
    "  ///       model was trained on: `batchSize * imageWidth * imageHeight * componentsCount`.\n",
    "  ///   - isModelQuantized: Whether the model is quantized (i.e. fixed point values rather than\n",
    "  ///       floating point values).\n",
    "  /// - Returns: The RGB data representation of the image buffer or `nil` if the buffer could not be\n",
    "  ///     converted.\n",
    "  private func rgbDataFromBuffer(\n",
    "    _ buffer: CVPixelBuffer,\n",
    "    byteCount: Int\n",
    "    ) -> Data? {\n",
    "    CVPixelBufferLockBaseAddress(buffer, .readOnly)\n",
    "    defer { CVPixelBufferUnlockBaseAddress(buffer, .readOnly) }\n",
    "    guard let mutableRawPointer = CVPixelBufferGetBaseAddress(buffer) else {\n",
    "      return nil\n",
    "    }\n",
    "    let count = CVPixelBufferGetDataSize(buffer)\n",
    "    let bufferData = Data(bytesNoCopy: mutableRawPointer, count: count, deallocator: .none)\n",
    "    var rgbBytes = [Float](repeating: 0, count: byteCount)\n",
    "    var index = 0\n",
    "    for component in bufferData.enumerated() {\n",
    "      let offset = component.offset\n",
    "      let isAlphaComponent = (offset % alphaComponent.baseOffset) == alphaComponent.moduloRemainder\n",
    "      guard !isAlphaComponent else { continue }\n",
    "      rgbBytes[index] = Float(component.element) / 255.0\n",
    "      index += 1\n",
    "    }\n",
    "\n",
    "    return rgbBytes.withUnsafeBufferPointer(Data.init)\n",
    "\n",
    "  }\n",
    "}\n",
    "\n",
    "// MARK: - Extensions\n",
    "\n",
    "extension Data {\n",
    "  /// Creates a new buffer by copying the buffer pointer of the given array.\n",
    "  ///\n",
    "  /// - Warning: The given array's element type `T` must be trivial in that it can be copied bit\n",
    "  ///     for bit with no indirection or reference-counting operations; otherwise, reinterpreting\n",
    "  ///     data from the resulting buffer has undefined behavior.\n",
    "  /// - Parameter array: An array with elements of type `T`.\n",
    "  init<T>(copyingBufferOf array: [T]) {\n",
    "    self = array.withUnsafeBufferPointer(Data.init)\n",
    "  }\n",
    "}\n",
    "\n",
    "extension Array {\n",
    "  /// Creates a new array from the bytes of the given unsafe data.\n",
    "  ///\n",
    "  /// - Warning: The array's `Element` type must be trivial in that it can be copied bit for bit\n",
    "  ///     with no indirection or reference-counting operations; otherwise, copying the raw bytes in\n",
    "  ///     the `unsafeData`'s buffer to a new array returns an unsafe copy.\n",
    "  /// - Note: Returns `nil` if `unsafeData.count` is not a multiple of\n",
    "  ///     `MemoryLayout<Element>.stride`.\n",
    "  /// - Parameter unsafeData: The data containing the bytes to turn into an array.\n",
    "  init?(unsafeData: Data) {\n",
    "\n",
    "    guard unsafeData.count % MemoryLayout<Element>.stride == 0 else { return nil }\n",
    "    #if swift(>=5.0)\n",
    "    self = unsafeData.withUnsafeBytes { .init($0.bindMemory(to: Element.self)) }\n",
    "    #else\n",
    "    self = unsafeData.withUnsafeBytes {\n",
    "      .init(UnsafeBufferPointer<Element>(\n",
    "        start: $0,\n",
    "        count: unsafeData.count / MemoryLayout<Element>.stride\n",
    "      ))\n",
    "    }\n",
    "    #endif  // swift(>=5.0)\n",
    "  }\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff3c49c7-a4d3-4eca-8bf0-4e0cdbd337d0",
   "metadata": {
    "height": 132
   },
   "outputs": [],
   "source": [
    "prompt_template = \"\"\"\n",
    "Can you please explain how this code works?\n",
    "\n",
    "{question}\n",
    "\n",
    "Use a lot of detail and make it as clear as possible.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90c20db9-f6e9-4d27-96b5-d43f0b24a484",
   "metadata": {
    "height": 81
   },
   "outputs": [],
   "source": [
    "completion = generate_text(\n",
    "    prompt = prompt_template.format(question=CODE_BLOCK)\n",
    ")\n",
    "print(completion.result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dcb1824-56dd-4e34-8079-9767cd9539e0",
   "metadata": {},
   "source": [
    "#### Try it out on your own code!\n",
    "- Try inputting some code into the `CODE_BLOCK` variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42ddae3b-5936-4b8e-b5ac-d6671c7a03e6",
   "metadata": {
    "height": 336
   },
   "outputs": [],
   "source": [
    "CODE_BLOCK = \"\"\" \n",
    "# replace this with your own code\n",
    "def foo(a):\n",
    "  b = a + 1\n",
    "  return 2*b\n",
    "\"\"\"\n",
    "\n",
    "prompt_template = \"\"\"\n",
    "Can you please explain how this code works?\n",
    "\n",
    "{question}\n",
    "\n",
    "Use a lot of detail and make it as clear as possible.\n",
    "\"\"\"\n",
    "\n",
    "completion = generate_text(\n",
    "    prompt = prompt_template.format(question=CODE_BLOCK)\n",
    ")\n",
    "print(completion.result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e93864a7-fd47-427b-b6f1-b492e159560d",
   "metadata": {},
   "source": [
    "### Ask an LLM to document a complex code base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57ac2cb3-0865-4d9b-89ba-bfc237ebad3b",
   "metadata": {
    "height": 149
   },
   "outputs": [],
   "source": [
    "prompt_template = \"\"\"\n",
    "Please write technical documentation for this code and \\n\n",
    "make it easy for a non swift developer to understand:\n",
    "\n",
    "{question}\n",
    "\n",
    "Output the results in markdown\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dcb82c6-a870-476b-b587-688d0344f33d",
   "metadata": {
    "height": 81
   },
   "outputs": [],
   "source": [
    "completion = generate_text(\n",
    "    prompt = prompt_template.format(question=CODE_BLOCK)\n",
    ")\n",
    "print(completion.result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9676f08-dd6a-4e83-9022-eb14ce4775f9",
   "metadata": {},
   "source": [
    "### Try it out on your own code!\n",
    "- Notice that we've modified the prompt slightly to refer to Python instead of Swift."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17699e40-d85e-4b9e-a569-e55dcdc2a23b",
   "metadata": {
    "height": 353
   },
   "outputs": [],
   "source": [
    "CODE_BLOCK = \"\"\" \n",
    "# replace this with your own code\n",
    "def foo(a):\n",
    "  b = a + 1\n",
    "  return 2*b\n",
    "\"\"\"\n",
    "\n",
    "prompt_template = \"\"\"\n",
    "Please write technical documentation for this code and \\n\n",
    "make it easy for a non Python developer to understand:\n",
    "\n",
    "{question}\n",
    "\n",
    "Output the results in markdown\n",
    "\"\"\"\n",
    "\n",
    "completion = generate_text(\n",
    "    prompt = prompt_template.format(question=CODE_BLOCK)\n",
    ")\n",
    "print(completion.result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcd7fd23-141c-477a-97c9-e7723df9d424",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
